use axum::{
    async_trait,
    extract::FromRequestParts,
    headers::{authorization::Bearer, Authorization},
    http::request::Parts,
    RequestPartsExt, TypedHeader,
};
use jsonwebtoken as jwt;
use jwt::{DecodingKey, EncodingKey, Validation};
use once_cell::sync::Lazy;
use redis::{aio::Connection, Client, RedisResult};
use serde::{Deserialize, Serialize};
use uuid::Uuid;

use crate::{error::ServerError, util::datetime::current_timestamp};

pub const JWT_TTL: usize = 24 * 60 * 60 * 1000;

#[derive(Debug, Serialize, Deserialize)]
pub struct Claims {
    pub user_id: Uuid,
    // 此字段名称不要改变，变了jwt认证会失败
    pub exp: usize,
}

impl Claims {
    pub fn new(user_id: Uuid) -> Self {
        Claims {
            user_id,
            // 有效期1天
            exp: current_timestamp() + JWT_TTL,
        }
    }

    pub fn get_jwt(&self) -> String {
        jwt::encode(&jwt::Header::default(), &self, &KEYS.encoding_key).unwrap()
    }
}

#[async_trait]
impl<S> FromRequestParts<S> for Claims
where
    S: Send + Sync,
{
    type Rejection = ServerError;

    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
        // Extract the token from the authorization header
        let TypedHeader(Authorization(bearer)) = parts
            .extract::<TypedHeader<Authorization<Bearer>>>()
            .await
            .map_err(|_| ServerError::Authentication)?;
        let token_data =
            jwt::decode::<Claims>(bearer.token(), &KEYS.decoding_key, &Validation::default())
                .map_err(|_| ServerError::Authentication)?;
        if let Ok(true) = token_check(&token_data.claims).await {
            Ok(token_data.claims)
        } else {
            Err(ServerError::Authentication)
        }
    }
}

/// token是否失效检查
async fn token_check(claims: &Claims) -> RedisResult<bool> {
    let client = Client::open("redis://127.0.0.1/")?;
    let mut con = client.get_async_connection().await?;
    // 注销登录检查
    if let 1 = redis::cmd("EXISTS")
        .arg(claims.get_jwt())
        .query_async::<Connection, usize>(&mut con)
        .await?
    {
        return Ok(false);
    }
    // 密码修改检查
    if let Some(ts) = redis::cmd("GET")
        .arg(claims.user_id.to_string())
        .query_async::<Connection, Option<usize>>(&mut con)
        .await?
    {
        if ts > claims.exp - JWT_TTL {
            return Ok(false);
        }
    }
    Ok(true)
}

/// secret key
static KEYS: Lazy<Keys> = Lazy::new(|| {
    let secret = std::env::var("JWT_SECRET").unwrap_or("axum-jwt".to_string());
    Keys::new(secret.as_str())
});

struct Keys {
    encoding_key: EncodingKey,
    decoding_key: DecodingKey,
}

impl Keys {
    pub fn new(secret: &str) -> Self {
        let bytes = secret.as_bytes();
        let encoding_key = jwt::EncodingKey::from_secret(bytes);
        let decoding_key = jwt::DecodingKey::from_secret(bytes);
        Keys {
            encoding_key,
            decoding_key,
        }
    }
}
